import asyncio
import base64
import json
import os

import uvicorn
from aiortc import RTCPeerConnection, RTCSessionDescription, MediaStreamTrack, sdp
from fastapi import FastAPI, Request
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from sse_starlette.sse import EventSourceResponse

from config import config, application_path
from video_receiver import VideoReceiver

app = FastAPI()

static = application_path + "/static"
if not os.path.exists(static) or not os.path.isdir(static):
    os.mkdir(static)

peer_connections = {}
message_queue = asyncio.Queue()


class Offer(BaseModel):
    sdp: str
    type: str


@app.get('/webrtc/createOffer')
async def create_offer(offer: str, request: Request):
    pc = RTCPeerConnection()

    @pc.on("track")
    def on_track(track):
        if isinstance(track, MediaStreamTrack):
            print(f"接收 {track.kind} 类型的轨道")
            if track.kind == 'video':
                video_receiver = VideoReceiver(message_queue)
                asyncio.ensure_future(video_receiver.handle_track(track))

    @pc.on("datachannel")
    def on_datachannel(channel):
        print(f"数据通道 {channel.label} 已建立")

    @pc.on("connectionstatechange")
    async def on_connectionstatechange():
        print(f"连接状态是 {pc.connectionState}")
        if pc.connectionState == "connected":
            print("WebRTC 连接成功建立")
        elif pc.connectionState == "closed":
            print("WebRTC 连接断开")
            del peer_connections[pc.cname]
            await pc.close()

    # Parse the offer
    offer_obj = json.loads(offer)
    offer_sdp = RTCSessionDescription(sdp=offer_obj['sdp'], type=offer_obj['type'])

    start_index = offer_sdp.sdp.find('cname:') + 6
    end_index = offer_sdp.sdp.find('\r', start_index)
    cname = offer_sdp.sdp[start_index:end_index]
    cname = base64.b64encode(cname.encode('utf-8')).decode('utf-8')
    peer_connections[cname] = pc
    pc.cname = cname

    await pc.setRemoteDescription(offer_sdp)

    # Create an answer and set it as the local description
    answer = await pc.createAnswer()
    await pc.setLocalDescription(answer)

    # Send the answer back to the client
    await message_queue.put({
        "resultCode": 200,
        "data": {
            "type": "answer",
            "data": answer.sdp
        }
    })

    async def event_generator():
        while True:
            if await request.is_disconnected():
                break
            # 测试取随机数据，每次取一个随机数
            message = await message_queue.get()
            yield {'data': json.dumps(message)}
            message_queue.task_done()

    return EventSourceResponse(event_generator())


class IceCandidate(BaseModel):
    candidate: str
    sdpMLineIndex: int
    sdpMid: str


@app.post('/webrtc/onIceCandidate')
async def on_ice_candidate(cname: str, body: IceCandidate):
    if not cname:
        return {'resultCode': 409, 'message': 'cname not provided'}

    ice_candidate = sdp.candidate_from_sdp(body.candidate)
    ice_candidate.sdpMid = body.sdpMid
    ice_candidate.sdpMLineIndex = body.sdpMLineIndex

    pc = peer_connections[cname]
    if not pc:
        return {'resultCode': 409, 'message': 'cname not found' + cname}

    await pc.addIceCandidate(ice_candidate)
    return {'resultCode': 200}


if __name__ == '__main__':
    port = config.getint('DEFAULT', 'port', fallback=10057)
    app.mount("/", StaticFiles(directory=static), name="static")
    uvicorn.run(app, host="0.0.0.0", port=port)
